"""
torch.cat   合并张量
torch.cat([x_1, h_0], 1)  再1 轴上合并数据
x_1 --> 1,3 
h_0 --> 1,4
合并后    1,7
"""
from torch import nn
import torch

import torch.nn.functional as F

torch.manual_seed(1)


def forward(x_1, h_0, c_0, w_ii, b_ii, w_hi, b_hi, w_if, b_if, w_hf, b_hf, w_io, b_io, w_ho, b_ho, w_ig, b_ig, w_hg,
            b_hg):
    # input gate
    i_1 = torch.sigmoid(torch.mm(x_1, w_ii) + b_ii + torch.mm(h_0, w_hi) + b_hi)
    # forget gate
    f_1 = torch.sigmoid(torch.mm(x_1, w_if) + b_if + torch.mm(h_0, w_hf) + b_hf)
    # output gate
    o_1 = torch.sigmoid(torch.mm(x_1, w_io) + b_io + torch.mm(h_0, w_ho) + b_ho)

    g_1 = torch.tanh(torch.mm(x_1, w_ig) + b_ig + torch.mm(h_0, w_hg) + b_hg)
    #
    c_1 = f_1 * c_0 + i_1 * g_1
    h_1 = o_1 * torch.tanh(c_1)

    return c_1, h_1


def forward_new(x_1, h_0, c_0, w_ii, b_ii, w_hi, b_hi, w_if, b_if, w_hf, b_hf, w_io, b_io, w_ho, b_ho, w_ig, b_ig, w_hg,
                b_hg):
    x_h = torch.cat([x_1, h_0], 1)
    print(x_h.shape)
    # input gate
    w_i = torch.cat([w_ii, w_hi])
    b_i = b_ii + b_hi
    i_1 = torch.sigmoid(torch.mm(x_h, w_i) + b_i)

    # forget gate
    f_1 = torch.sigmoid(torch.mm(x_1, w_if) + b_if + torch.mm(h_0, w_hf) + b_hf)
    # output gate
    o_1 = torch.sigmoid(torch.mm(x_1, w_io) + b_io + torch.mm(h_0, w_ho) + b_ho)

    g_1 = torch.tanh(torch.mm(x_1, w_ig) + b_ig + torch.mm(h_0, w_hg) + b_hg)
    #
    c_1 = f_1 * c_0 + i_1 * g_1
    h_1 = o_1 * torch.tanh(c_1)

    return c_1, h_1


word_dimension = 5  # 每个单词转化成5维向量
h_dimension = 6  # 隐藏层的维度是6

# input gate
w_ii = torch.randn(word_dimension, 1)
b_ii = torch.randn(1, 1)
w_hi = torch.randn(h_dimension, 1)
b_hi = torch.randn(1, 1)

# forget gate
w_if = torch.randn(word_dimension, 1)
b_if = torch.randn(1, 1)
w_hf = torch.randn(h_dimension, 1)
b_hf = torch.randn(1, 1)

# output gate
w_io = torch.randn(word_dimension, 1)
b_io = torch.randn(1, 1)
w_ho = torch.randn(h_dimension, 1)
b_ho = torch.randn(1, 1)

w_ig = torch.randn(word_dimension, h_dimension)
b_ig = torch.randn(1, h_dimension)
w_hg = torch.randn(h_dimension, h_dimension)
b_hg = torch.randn(1, h_dimension)

torch.manual_seed(100)
arrive = torch.randn(1, 5)
on = torch.randn(1, 5)
taibei = torch.randn(1, 5)

h_0 = torch.zeros(1, h_dimension)
c_0 = torch.zeros(1, h_dimension)

c_1, h_1 = forward(arrive, h_0, c_0, w_ii, b_ii, w_hi, b_hi, w_if, b_if, w_hf, b_hf, w_io, b_io, w_ho, b_ho, w_ig, b_ig,
                   w_hg,
                   b_hg)

print(c_1)
print(h_1)

c_1, h_1 = forward_new(arrive, h_0, c_0, w_ii, b_ii, w_hi, b_hi, w_if, b_if, w_hf, b_hf, w_io, b_io, w_ho, b_ho, w_ig,
                       b_ig,
                       w_hg,
                       b_hg)
print(c_1)
print(h_1)

c_2, h_2 = forward_new(on, h_0, c_0, w_ii, b_ii, w_hi, b_hi, w_if, b_if, w_hf, b_hf, w_io, b_io, w_ho, b_ho, w_ig,
                      b_ig,
                      w_hg,
                      b_hg)

print(c_2)
print(h_2)
